{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "from easydict import EasyDict as edict\n",
    "import matplotlib.pyplot as plt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA = edict()\n",
    "DATA.grid = edict()\n",
    "DATA.grid.shape = (7, 10)\n",
    "DATA.grid.startPoint = (3, 0)\n",
    "DATA.grid.endPoint = (3, 7)\n",
    "DATA.grid.wind = (0 ,0, 0, 1, 1, 1, 2, 2, 1, 0)\n",
    "DATA.actionSet = [(0, 1), (1, 0), (0, -1), (-1, 0), (1, 1), (1, -1), (-1, -1), (-1, 1), (0, 0)]\n",
    "DATA.episode = edict({'S':[], 'A':[], 'R':[]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "def drawPictures(state, order, num):\n",
    "\n",
    "    img = np.zeros((30 * 7, 30 * 10, 3), np.uint8) #生成一个空灰度图像\n",
    "    img.fill(255)\n",
    "    # wallColor = (0, 0, 0)    # BGR 黑色\n",
    "    trackColor = (0, 255, 0) # 绿色\n",
    "    startColor = (255, 0, 0)\n",
    "    endColor = (0, 0, 255)\n",
    "    stateColor = (128, 0, 128)\n",
    "    \n",
    "\n",
    "    for i in range(7):\n",
    "        for j in range(10):\n",
    "            cv2.rectangle(img, (30 * j, 30 * i), (30 * (j + 1), 30 * (i + 1)), trackColor, thickness=2)\n",
    "\n",
    "    cv2.rectangle(img, (30 * 0, 30 * 3), (30 * (0 + 1), 30 * (3 + 1)), startColor, thickness=-1)\n",
    "    cv2.rectangle(img, (30 * 7, 30 * 3), (30 * (7 + 1), 30 * (3 + 1)), endColor, thickness=-1)\n",
    "    cv2.rectangle(img, (30 * state[1], 30 * state[0]), (30 * (state[1] + 1), 30 * (state[0] + 1)), stateColor, thickness=-1)\n",
    "    cv2.imwrite(\"./images\" + str(num) + \"/step\" + str(order) + \".jpg\", img)\n",
    "\n",
    "def getVideo(num):\n",
    "    import cv2\n",
    "    import os\n",
    "    \n",
    "    # 读取时序图中的第一张图片\n",
    "    img = cv2.imread('./images' + str(num) + '/step0.jpg')\n",
    "    # 设置每秒读取多少张图片\n",
    "    fps = 2\n",
    "    imgInfo = img.shape\n",
    "    \n",
    "    # 获取图片宽高度信息\n",
    "    size = (imgInfo[1], imgInfo[0])\n",
    "    fourcc = cv2.VideoWriter_fourcc(*\"MPEG\")\n",
    "    \n",
    "    # 定义写入图片的策略\n",
    "    videoWrite = cv2.VideoWriter('output' + str(num) + '.mp4', fourcc, fps, size)\n",
    "    \n",
    "    files = os.listdir('./images' + str(num))\n",
    "\n",
    "    out_num = len(files)\n",
    "    for i in range(0, out_num):\n",
    "        # 读取所有的图片\n",
    "        fileName = './images' + str(num) + '/step' + str(i) + '.jpg'\n",
    "        img = cv2.imread(fileName)\n",
    "        # print(out_num)\n",
    "        # 将图片写入所创建的视频对象\n",
    "        videoWrite.write(img)\n",
    "    \n",
    "    videoWrite.release()\n",
    "    print('finish')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "LEGAL_ACTION = [[[] for j in range(10)] for i in range(7)]\n",
    "def getPossiableActions(state, num_action):\n",
    "    if LEGAL_ACTION[state[0]][state[1]] != []:\n",
    "        return LEGAL_ACTION[state[0]][state[1]]\n",
    "    else:\n",
    "        for i in range(num_action):\n",
    "            if state[0] + DATA.actionSet[i][0] <= 6 and state[0] + DATA.actionSet[i][0] >= 0 and state[1] + DATA.actionSet[i][1] <= 9 and state[1] + DATA.actionSet[i][1] >= 0:\n",
    "                LEGAL_ACTION[state[0]][state[1]].append(DATA.actionSet[i])\n",
    "        return LEGAL_ACTION[state[0]][state[1]]\n",
    "\n",
    "\n",
    "def actionMapTo2D(num):\n",
    "    return DATA.actionSet[num]\n",
    "\n",
    "def actionMapTo1D(action2D, num_action):\n",
    "    for i in range(num_action):\n",
    "        if action2D == DATA.actionSet[i]:\n",
    "            return i\n",
    "\n",
    "\n",
    "def getBestPolicy(Q, num_action):\n",
    "    policy = np.zeros(shape = (7, 10), dtype = int)\n",
    "    for i in range(7):\n",
    "        for j in range(10):\n",
    "            state = (i, j)\n",
    "            legalAction = getPossiableActions(state, num_action)\n",
    "            maxAction = legalAction[0]\n",
    "            maxAction1D = actionMapTo1D(maxAction, num_action)\n",
    "            maxActionValue = Q[i][j][maxAction1D]\n",
    "            length = len(legalAction)\n",
    "            for k in range(1, length):\n",
    "                action1D = actionMapTo1D(legalAction[k], num_action)\n",
    "                if Q[i][j][action1D] > maxActionValue:\n",
    "                    maxAction = legalAction[k]\n",
    "                    maxAction1D = action1D\n",
    "                    maxActionValue = Q[i][j][action1D]\n",
    "            policy[i][j] = maxAction1D\n",
    "    return policy\n",
    "\n",
    "def getBestTrace(Q, num_action):\n",
    "    policy = getBestPolicy(Q, num_action)\n",
    "    state = np.array(DATA.grid.startPoint, dtype = int)\n",
    "    endState = np.array(DATA.grid.endPoint, dtype = int)\n",
    "    bestTrace = [state]\n",
    "    while True:\n",
    "        action1D = policy[state[0]][state[1]]\n",
    "        action2D = actionMapTo2D(action1D)\n",
    "\n",
    "        if (state == endState).all() == True:\n",
    "            break\n",
    "        else:\n",
    "\n",
    "            state = np.array([min(max(state[0] + action2D[0] - DATA.grid.wind[state[1]], 0), 6), state[1] + action2D[1]], dtype = int)\n",
    "        # if data.track[state[0]][state[1]] == -1:\n",
    "        #     state = env.start()\n",
    "        print(state)\n",
    "        bestTrace.append(state)\n",
    "    return bestTrace\n",
    "\n",
    "def visualizer(Q, num_action):\n",
    "    bestTrace = getBestTrace(Q, num_action)\n",
    "    length = len(bestTrace)\n",
    "    for i in range(length):\n",
    "        drawPictures(bestTrace[i], i, num_action)\n",
    "    getVideo(num_action)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[3 1]\n",
      "[3 2]\n",
      "[3 3]\n",
      "[1 3]\n",
      "[0 4]\n",
      "[0 5]\n",
      "[0 6]\n",
      "[0 7]\n",
      "[0 8]\n",
      "[0 9]\n",
      "[1 9]\n",
      "[2 9]\n",
      "[3 9]\n",
      "[4 9]\n",
      "[4 8]\n",
      "[3 7]\n",
      "finish\n",
      "[3 1]\n",
      "[4 2]\n",
      "[5 3]\n",
      "[4 4]\n",
      "[4 5]\n",
      "[4 6]\n",
      "[3 7]\n",
      "finish\n",
      "[3 1]\n",
      "[4 2]\n",
      "[5 3]\n",
      "[5 4]\n",
      "[4 5]\n",
      "[4 6]\n",
      "[3 7]\n",
      "finish\n"
     ]
    }
   ],
   "source": [
    "Q4 = np.load(\"Q4_6.9.npy\")\n",
    "Q8 = np.load(\"Q8_6.9.npy\")\n",
    "Q9 = np.load(\"Q9_6.9.npy\")\n",
    "\n",
    "visualizer(Q4, 4)\n",
    "LEGAL_ACTION = [[[] for j in range(10)] for i in range(7)]\n",
    "visualizer(Q8, 8)\n",
    "LEGAL_ACTION = [[[] for j in range(10)] for i in range(7)]\n",
    "visualizer(Q9, 9)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "494899efd6527d56ea7f55c588d0081523a17dc3a9ff1107f3394ad815ff2527"
  },
  "kernelspec": {
   "display_name": "Python 3.7.7 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
